/* Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

        http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
        limitations under the License.
==============================================================================*/

#include "emb_table/embedding_table.h"
#include "log/logger.h"
#include "utils/singleton.h"
#include "hd_transfer/hd_transfer.h"
#include "file_system/file_system_handler.h"

using namespace MxRec;

EmbeddingTable::EmbeddingTable()
{
}

EmbeddingTable::EmbeddingTable(const EmbInfo& info, const RankInfo& rankInfo, int inSeed)
    : name(info.name), hostVocabSize(info.hostVocabSize), devVocabSize(info.devVocabSize),
      ssdVocabSize(info.ssdVocabSize), freeSize_(0), maxOffset(0), isDynamic_(rankInfo.useDynamicExpansion),
      embSize_(info.embeddingSize), extEmbSize_(info.extEmbeddingSize),
      embInfo_(info), seed_(inSeed), rankId_(rankInfo.rankId), rankSize_(rankInfo.rankSize)
{
    LOG_INFO("Init embedding table:{}, isDynamic:{}, embeddingSize:{}, extEmbeddingSize:{}.",
             name, isDynamic_, embSize_, extEmbSize_);
}

EmbeddingTable::~EmbeddingTable()
{
}

void EmbeddingTable::Key2Offset(std::vector<emb_key_t>& keys, int channel)
{
    return;
}

void EmbeddingTable::Key2OffsetForDp(std::vector<emb_key_t>& keys, int channel)
{
}

size_t EmbeddingTable::GetMaxOffset() const
{
    return maxOffset;
}

int64_t EmbeddingTable::capacity() const
{
    return static_cast<int64_t>(devVocabSize);
}

size_t EmbeddingTable::size() const
{
    return maxOffset;
}

void EmbeddingTable::EvictKeys(const std::vector<emb_cache_key_t>& keys)
{
    std::unique_lock<std::shared_mutex> lock(keyOffsetMutex_); // lock for PROCESS_THREAD
    size_t keySize = keys.size();
    for (size_t i = 0; i < keySize; i++) {
        emb_key_t key = keys[i];
        if (key == INVALID_KEY_VALUE) {
            LOG_WARN("Evict key is INVALID_KEY_VALUE, table:{}.", name);
            continue;
        }
        const auto& iter = keyOffsetMap.find(key);
        if (iter == keyOffsetMap.end()) { // not found
            continue;
        }
        keyOffsetMap.erase(iter);
        evictDevPos.emplace_back(iter->second);
        LOG_TRACE("EvictKeys table:{}, offset:{}.", name, iter->second);
    }
    LOG_INFO("EvictKeys table:{}, evict key size:{}.", name, evictDevPos.size());
}

const std::vector<int64_t>& EmbeddingTable::GetEvictedKeys()
{
    return evictDevPos;
}

const std::vector<int64_t>& EmbeddingTable::GetHostEvictedKeys()
{
    return evictHostPos;
}

void EmbeddingTable::EvictInitDeviceEmb()
{
    if (evictDevPos.size() > devVocabSize) {
        auto errMsg = Logger::Format(
            "Overflow! Init evicted embedding on device, table:{}, evictOffset size:{} bigger than devVocabSize:{}.",
            name, evictDevPos.size(), devVocabSize);
        auto error = Error(ModuleName::M_EMB_TABLE, ErrorType::LOGIC_ERROR, errMsg);
        LOG_ERROR(error.ToString());
        throw std::runtime_error(error.ToString());
    }

    vector<Tensor> tmpDataOut;
    Tensor tmpData = Vec2TensorI32(evictDevPos);
    tmpDataOut.emplace_back(tmpData);
    tmpDataOut.emplace_back(Tensor(tensorflow::DT_INT32, { 1 }));

    auto evictLen = tmpDataOut.back().flat<int32>();
    evictLen(0) = static_cast<int>(evictDevPos.size());

    // evict key发送给dev侧，dev侧初始化emb
    auto trans = Singleton<HDTransfer>::GetInstance();
    trans->Send(TransferChannel::EVICT, tmpDataOut, TRAIN_CHANNEL_ID, name);

    LOG_INFO("EvictInitDeviceEmb, table:{}, send evict offsetSize:{}.", name, evictDevPos.size());
}

absl::flat_hash_map<emb_key_t, int64_t> EmbeddingTable::GetKeyOffsetMap()
{
    return keyOffsetMap;
}

void EmbeddingTable::SetFileSystemPtr(const string& savePath)
{
    unique_ptr<FileSystemHandler> fileSystemHandler = make_unique<FileSystemHandler>();
    fileSystemPtr_ = fileSystemHandler->Create(savePath);
}

void EmbeddingTable::UnsetFileSystemPtr()
{
    fileSystemPtr_ = nullptr;
}

vector<int64_t> EmbeddingTable::GetLoadOffset()
{
    return loadOffset;
}

void EmbeddingTable::Load(const string& savePath, map<string, unordered_set<emb_cache_key_t>>& trainKeySet,
                          const vector<string>& warmStartTables)
{
}

void EmbeddingTable::Save(const string& savePath, const int pythonBatchId, bool saveDelta,
                          const map<emb_key_t, KeyInfo>& keyInfo)
{
}

void EmbeddingTable::BackUpTrainStatus()
{
}

void EmbeddingTable::RecoverTrainStatus()
{
}

void EmbeddingTable::MakeDir(const string& dirName)
{
    CheckFileSystemPtr();
    fileSystemPtr_->CreateDir(dirName);
}

void EmbeddingTable::SetCacheManager(CacheManager *cacheManager)
{
}

TableInfo EmbeddingTable::GetTableInfo()
{
    TableInfo ti = {
        .name=name,
        .hostVocabSize=hostVocabSize,
        .devVocabSize=devVocabSize,
        .maxOffset=maxOffset,
        .keyOffsetMap=keyOffsetMap,
    };
    return ti;
}

vector<int64_t> EmbeddingTable::GetDeviceOffset()
{
    return vector<int64_t>{};
}

void EmbeddingTable::SetOptimizerInfo(OptimizerInfo& optimizerInfo)
{
}

void EmbeddingTable::SetHDTransfer(HDTransfer *hdTransfer)
{
}

void EmbeddingTable::SetEmbCache(ock::ctr::EmbCacheManagerPtr embCache)
{
}

void EmbeddingTable::CheckFileSystemPtr() const
{
    if (fileSystemPtr_ != nullptr) {
        return;
    }
    auto error = Error(ModuleName::M_EMB_TABLE, ErrorType::NULL_PTR,
                       StringFormat("Failed to obtain the file system pointer,"
                                    " the file system pointer is null, table:%s.", name.c_str()));
    LOG_ERROR(error.ToString());
    throw std::runtime_error(error.ToString());
}

void EmbeddingTable::CheckReadKeyFileSize(const string& fileName, size_t fileSize)
{
    if (fileSize < FILE_MAX_SIZE) {
        return;
    }

    auto error = Error(ModuleName::M_EMB_TABLE, ErrorType::LOGIC_ERROR,
                       StringFormat("Load keys failed, file %s size %d is too big, table:%s.",
                                    fileName.c_str(), fileSize, name.c_str()));
    LOG_ERROR(error.ToString());
    throw std::runtime_error(error.ToString());
}

void EmbeddingTable::CheckLoadKeyMallocPtr(const int64_t* mallocPtr, size_t mallocByteSize)
{
    if (mallocPtr != nullptr) {
        return;
    }
    auto error = Error(ModuleName::M_EMB_TABLE, ErrorType::LOGIC_ERROR,
                       StringFormat("Load keys failed. Failed to allocate %d bytes using malloc, table:%s.",
                                    mallocByteSize, name.c_str()));
    LOG_ERROR(error.ToString());
    throw std::runtime_error(error.ToString());
}

void EmbeddingTable::CheckReadKeyFileBytes(ssize_t readReturnCode, const string& fileName, size_t fileSize)
{
    if (readReturnCode == -1) {
        auto error = Error(ModuleName::M_EMB_TABLE, ErrorType::LOGIC_ERROR,
                           StringFormat("Load keys failed. An error occurred while reading file:%s, table:%s.",
                                        fileName.c_str(), name.c_str()));
        LOG_ERROR(error.ToString());
        throw std::runtime_error(error.ToString());
    }
    if (readReturnCode != fileSize) {
        string errMsg = StringFormat(
            "Load keys failed. Expected to read %d bytes, but actually read %d bytes to file %s, table:%s.",
            fileSize, readReturnCode, fileName.c_str(), name.c_str());
        auto error = Error(ModuleName::M_EMB_TABLE, ErrorType::LOGIC_ERROR, errMsg);
        LOG_ERROR(error.ToString());
        throw std::runtime_error(error.ToString());
    }
}

unordered_set<int64_t> EmbeddingTable::GetPaddingKeysOffset()
{
    return paddingKeysOffset;
}

void EmbeddingTable::RecordPaddingKeysOffset(int channel, emb_key_t key, int64_t offset)
{
    if (channel != TRAIN_CHANNEL_ID || !embInfo_.paddingKeysMask) {
        return;
    }

    auto it = std::find(embInfo_.paddingKeys.begin(), embInfo_.paddingKeys.end(), key);
    if (it != embInfo_.paddingKeys.end()) {
        paddingKeysOffset.insert(offset);
    }
}

void EmbeddingTable::SyncLatestEmbedding(int pythonBatchId)
{
}
